#include <tinx/interrupt.h>
#include <tinx/debug.h>
#include <tinx/device.h>
#include <tinx/task.h>
#include <tinx/memory.h>
#include <tinx/ide.h>
#include <tinx/part.h>
#include <tinx/pci.h>
#include <tinx/io.h>
#include <tinx/arena.h>
#include <tinx/vsprintf.h>
#include <tinx/string.h>
#include <tinx/assert.h>
#include <tinx/stdlib.h>

#define IDE1_IOBASE 0x1F0
#define IDE2_IOBASE 0x170

#define IDE_CTRL_NR 1
#define IDE_CTRL_DISK_NR 2

#define IDE_PIO_READ_CMD 0x20   // PIO 读硬盘命令
#define IDE_PIO_WRITE_CMD 0x30  // PIO 写硬盘命令
#define IDE_PIO_IDENTIFY_CMD 0xEC

#define IDE_UDMA_READ_CMD 0xC8   // UDMA 读硬盘命令
#define IDE_UDMA_WRITE_CMD 0xCA  // UDMA 写硬盘命令

#define IDE_DIAGNOSTIC_CMD 0x90  // 诊断命令

#define IDE_LAST_PRD 0x8000

#define IDE_DATA_REG 0
#define IDE_ERROR_REG 1
#define IDE_SECTORS_REG 2
#define IDE_LBA_LOW_REG 3
#define IDE_LBA_MID_REG 4
#define IDE_LBA_HIGH_REG 5
#define IDE_HDDEVSEL_REG 6
#define IDE_COMMAND_REG 7
#define IDE_STATUS_REG 7

#define IDE_NULL 0b00000000
#define IDE_DRQ  0b00001000 // 硬盘数据准备完毕
#define IDE_ERR  0b00000001 // 硬盘出错
#define IDE_BSY  0b10000000 // 硬盘繁忙
#define IDE_RDY  0b01000000 // 硬盘就绪

// 总线主控寄存器偏移
#define BM_COMMAND_REG 0 // 命令寄存器偏移
#define BM_STATUS_REG 2  // 状态寄存器偏移
#define BM_PRD_ADDR 4    // PRD 地址寄存器偏移

// 总线主控命令寄存器
#define BM_CR_STOP 0x00  // 终止传输
#define BM_CR_START 0x01 // 开始传输
#define BM_CR_WRITE 0x00 // 主控写磁盘
#define BM_CR_READ 0x08  // 主控读磁盘

// 总线主控状态寄存器
#define BM_SR_ACT 0x01     // 激活
#define BM_SR_ERR 0x02     // 错误
#define BM_SR_INT 0x04     // 中断信号生成
#define BM_SR_DRV0 0x20    // 驱动器 0 可以使用 DMA 方式
#define BM_SR_DRV1 0x40    // 驱动器 1 可以使用 DMA 方式
#define BM_SR_SIMPLEX 0x80 // 仅单纯形操作

enum IDE_INTERFACE
{
    IDE_INTERFACE_UNKNOW,
    IDE_INTERFACE_ATA,
    IDE_INTERFACE_ATAPI
};

typedef struct ide_params_t
{
    u16 config;                 // 0 General configuration bits
    u16 cylinders;              // 01 cylinders
    u16 RESERVED;               // 02
    u16 heads;                  // 03 heads
    u16 RESERVED[5 - 3];        // 05
    u16 sectors;                // 06 sectors per track
    u16 RESERVED[9 - 6];        // 09
    u8 serial[20];              // 10 ~ 19 序列号
    u16 RESERVED[22 - 19];      // 10 ~ 22
    u8 firmware[8];             // 23 ~ 26 固件版本
    u8 model[40];               // 27 ~ 46 模型数
    u8 drq_sectors;             // 47 扇区数量
    u8 RESERVED[3];             // 48
    u16 capabilities;           // 49 能力
    u16 RESERVED[59 - 49];      // 50 ~ 59
    u32 total_lba;              // 60 ~ 61
    u16 RESERVED;               // 62
    u16 mdma_mode;              // 63
    u8 RESERVED;                // 64
    u8 pio_mode;                // 64
    u16 RESERVED[79 - 64];      // 65 ~ 79 参见 ATA specification
    u16 major_version;          // 80 主版本
    u16 minor_version;          // 81 副版本
    u16 commmand_sets[87 - 81]; // 82 ~ 87 支持的命令集
    u16 RESERVED[118 - 87];     // 88 ~ 118
    u16 supiobase_settings;       // 119
    u16 enable_settings;        // 120
    u16 RESERVED[221 - 120];    // 221
    u16 transiobase_major;        // 222
    u16 transiobase_minor;        // 223
    u16 RESERVED[254 - 223];    // 254
    u16 integrity;              // 校验和
}_packed ide_params_t;

ide_ctrl_t ide_ctrls[IDE_CTRL_NR];

static int ide_busy_wait(ide_ctrl_t *ctrl, u8 mask)
{
    while (true)
    {
        // 从备用状态寄存器中读状态
        u8 state = inb(ctrl->iobase + IDE_STATUS_REG);
        if (state & IDE_ERR) // 有错误
        {
            u8 err = inb(ctrl->iobase + IDE_ERROR_REG);
            DEBUGK("IDE Error, error code: 0x%x\n", err);
        }
        if (state & IDE_BSY) // 驱动器忙
        {
            continue;
        }
        if ((state & mask) == mask)
        {
            return EOK;
        }
    }
}

static void ide_handler(u8 vector)
{
    send_eoi(vector);

    u8 irq = vector - 0x20;
    u8 idx = irq - IRQ_HARDDISK;

    ide_ctrl_t *ctrl = ide_ctrls + idx;
    u8 status = inb(ctrl->iobase + IDE_STATUS_REG);

    if (ctrl->waiter)
    {
        task_unblock(ctrl->waiter, 0);
        ctrl->waiter = NULL;
    }
}

static void ide_select_drive(ide_disk_t *disk)
{
    outb(disk->ctrl->iobase + IDE_HDDEVSEL_REG, disk->selector);
}

static void ide_send_cmd(ide_disk_t *disk, u8 cmd)
{
    outb(disk->ctrl->iobase + IDE_COMMAND_REG, cmd);
}

static void ide_select_sectors(ide_disk_t *disk, u32 lba, u8 count)
{
    u16 iobase = disk->ctrl->iobase;
    outb(iobase + IDE_SECTORS_REG, count);
    outb(iobase + IDE_LBA_LOW_REG, lba & 0xff);
    outb(iobase + IDE_LBA_MID_REG, (lba >> 8) & 0xff);
    outb(iobase + IDE_LBA_HIGH_REG, (lba >> 16) & 0xff);
    outb(iobase + IDE_HDDEVSEL_REG, disk->selector | ((lba >> 24) & 0xf));
}

void ide_pio_read_sector(ide_disk_t *disk, u16 *buf)
{
    u16 port = disk->ctrl->iobase + IDE_DATA_REG;
    for (size_t i = 0; i < SECTOR_SIZE / sizeof(u16); i++)
    {
        buf[i] = inw(port);
    }
}

void ide_pio_write_sector(ide_disk_t *disk, u16 *buf)
{
    u16 port = disk->ctrl->iobase + IDE_DATA_REG;
    for (size_t i = 0; i < SECTOR_SIZE / sizeof(u16); i++)
    {
        outw(port, buf[i]);
    }
}

int ide_pio_read(ide_disk_t *disk, void *addr, u8 count, u32 lba)
{
    if (count + lba > disk->sectors)
        return EOF;

    assert(!get_interrupt_state());
    assert(count != 0);

    task_t *task = running_task();

    // 获取对应的 IDE 通道
    ide_ctrl_t *ctrl = disk->ctrl;

    // 获得锁
    lock_acquire(&ctrl->lock);

    // 选择硬盘
    ide_select_drive(disk);

    // 等待硬盘就绪
    ide_busy_wait(ctrl, IDE_RDY);

    // 选择扇区
    ide_select_sectors(disk, lba, count);

    // 发送命令
    ide_send_cmd(disk, IDE_PIO_READ_CMD);

    for (size_t i = 0; i < count; i++)
    {
        // 等待硬盘
        ctrl->waiter = task;
        task_block(task, TASK_BLOCKED);
        ide_busy_wait(ctrl, IDE_DRQ);
        // 读取数据
        u32 offset = ((u32)addr + i * SECTOR_SIZE);
        ide_pio_read_sector(disk, (u16 *)offset);
    }

    // 释放锁
    lock_release(&ctrl->lock);
    return count;
}

int ide_pio_write(ide_disk_t *disk, void *addr, u8 count, u32 lba)
{
    if (count + lba > disk->sectors)
        return EOF;

    assert(!get_interrupt_state());
    assert(count != 0);

    task_t *task = running_task();

    // 获取对应的 IDE 通道
    ide_ctrl_t *ctrl = disk->ctrl;

    // 获得锁
    lock_acquire(&ctrl->lock);

    // 选择硬盘
    ide_select_drive(disk);

    // 等待硬盘就绪
    ide_busy_wait(ctrl, IDE_RDY);

    // 选择扇区
    ide_select_sectors(disk, lba, count);

    // 发送命令
    ide_send_cmd(disk, IDE_PIO_WRITE_CMD);

    for (size_t i = 0; i < count; i++)
    {
        // 写入数据
        u32 offset = ((u32)addr + i * SECTOR_SIZE);
        ide_pio_write_sector(disk, (u16 *)offset);
        // 等待硬盘
        ctrl->waiter = task;
        task_block(task, TASK_BLOCKED);
        ide_busy_wait(ctrl, IDE_DRQ);
    }

    // 释放锁
    lock_release(&ctrl->lock);
    return count;
}

int ide_ioctl(ide_disk_t *disk, int cmd, void *args, int flags)
{
    switch (cmd)
    {
    case DEV_CMD_SECTORS_COUNT:
        return disk->sectors;
    case DEV_CMD_SECTORS_START:
        return 0;
    }
    return EOF;
}

static bool ide_setup_udma(ide_ctrl_t *ctrl, u32 cmd, char *buf, u16 len)
{
    // 保证没有跨页
    if (!(((u32)buf + len) <= ((u32)buf & (~0xfff)) + PAGE_SIZE))
    {
        return false;
    }

    ide_prd_t *prd = &ctrl->prd;
    prd->paddr = get_paddr((u32)buf);
    prd->size = len;
    prd->flags = IDE_LAST_PRD;

    outl(ctrl->bmbase + BM_PRD_ADDR, get_paddr((u32)prd));

    outb(ctrl->bmbase + BM_COMMAND_REG, cmd | BM_CR_STOP);

    // 设置 UDMA 允许中断和错误
    u32 data = inb(ctrl->bmbase + BM_STATUS_REG) | BM_SR_INT | BM_SR_ERR;
    outb(ctrl->bmbase + BM_STATUS_REG, data);

    return true;
}

static void ide_start_udma(ide_ctrl_t *ctrl)
{
    u8 data = inb(ctrl->bmbase + BM_COMMAND_REG) | BM_CR_START;
    outb(ctrl->bmbase + BM_COMMAND_REG, data);
}

static void ide_stop_udma(ide_ctrl_t *ctrl)
{
    u8 data = inb(ctrl->bmbase + BM_COMMAND_REG);
    outb(ctrl->bmbase + BM_COMMAND_REG, data & (~BM_CR_START));

    u8 status = inb(ctrl->bmbase + BM_STATUS_REG);
    outb(ctrl->bmbase + BM_STATUS_REG, data | BM_SR_INT | BM_SR_ERR);

    assert((status & BM_SR_ERR) == 0);
}

int ide_udma_read(ide_disk_t *disk, void *addr, u8 count, u32 lba)
{
    if (count + lba > disk->sectors)
        return EOF;

    ide_ctrl_t *ctrl = disk->ctrl;

    lock_acquire(&ctrl->lock);

    // 设置 UDMA
    if (!ide_setup_udma(ctrl, BM_CR_READ, addr, count * SECTOR_SIZE))
    {
        return EOF;
    }

    // 选择扇区
    ide_select_sectors(disk, lba, count);

    // 发送读命令
    ide_send_cmd(disk, IDE_UDMA_READ_CMD);

    // 开始传输
    ide_start_udma(ctrl);

    // 堵塞当前任务
    task_t *task = running_task();
    ctrl->waiter = task;
    task_block(task, TASK_BLOCKED);

    // 停止传输
    ide_stop_udma(ctrl);

    lock_release(&ctrl->lock);

    return count;
}

int ide_udma_write(ide_disk_t *disk, void *addr, u8 count, u32 lba)
{
    if (count + lba > disk->sectors)
        return EOF;

    ide_ctrl_t *ctrl = disk->ctrl;

    lock_acquire(&ctrl->lock);

    // 设置 UDMA
    if (!ide_setup_udma(ctrl, BM_CR_WRITE, addr, count * SECTOR_SIZE))
    {
        return EOF;
    }

    // 选择扇区
    ide_select_sectors(disk, lba, count);

    // 发送写命令
    ide_send_cmd(disk, IDE_UDMA_WRITE_CMD);

    // 开始传输
    ide_start_udma(ctrl);

    // 堵塞当前任务
    task_t *task = running_task();
    ctrl->waiter = task;
    task_block(task, TASK_BLOCKED);

    // 停止传输
    ide_stop_udma(ctrl);

    lock_release(&ctrl->lock);

    return count;
}

static void ide_fixstrings(char *buf, u32 len)
{
    for (size_t i = 0; i < len; i += 2)
    {
        register char ch = buf[i];
        buf[i] = buf[i + 1];
        buf[i + 1] = ch;
    }
    buf[len - 1] = '\0';
}

bool ide_identify(ide_disk_t *disk)
{
    // 获得一些内存作为缓冲
    ide_params_t *parmas = kmalloc(SECTOR_SIZE);

    u16 iobase = disk->ctrl->iobase;

    outb(iobase + IDE_LBA_LOW_REG, 0);
    outb(iobase + IDE_LBA_MID_REG, 0);
    outb(iobase + IDE_LBA_HIGH_REG, 0);

    // 选择硬盘
    outb(iobase + IDE_HDDEVSEL_REG, disk->selector & 0b10110000);

    // 发送获得信息命令
    outb(disk->ctrl->iobase + IDE_COMMAND_REG, IDE_PIO_IDENTIFY_CMD);

    // 硬盘不存在状态为 0
    if (inb(disk->ctrl->iobase + IDE_STATUS_REG) == 0)
    {
        disk->ctrl = NULL;
        return false;
    }

    // 等待硬盘准备数据
    task_t *task = running_task();
    disk->ctrl->waiter = task;
    task_block(task, TASK_BLOCKED);

    ide_busy_wait(disk->ctrl, IDE_RDY);

    // 读信息
    ide_pio_read_sector(disk, (u16 *)parmas);

    disk->sectors = parmas->total_lba;

    kfree(parmas);

    return true;
}

void ide_device_init()
{
    ide_iotype_t iotype = IDE_IO_TYPE_PIO;
    u16 bmbase = 0;

    pci_device_t *device = pci_find_device_by_class(PCI_CLASS_STORAGE_IDE);
    if (device)
    {
        pci_bar_t bar;
        int ret = pci_find_bar(device, &bar, 4);
        assert(ret != EOF);

        // 打开总线主控
        pci_send_command(device, PCI_COMMAND_MASTER);

        iotype = IDE_IO_TYPE_UDMA;
        bmbase = bar.iobase;
    }

    for (size_t cidx = 0; cidx < IDE_CTRL_NR; cidx++)
    {
        ide_ctrl_t *ctrl = &ide_ctrls[cidx];
        ctrl->irq = 0x2E + cidx;
        ctrl->waiter = NULL;
        ctrl->iotype = iotype;
        ctrl->bmbase = bmbase;
        ctrl->iobase = cidx ? IDE2_IOBASE : IDE1_IOBASE;

        set_interrupt_handler(0x2e + cidx, ide_handler); // 注册中断
        set_interrupt_mask(0xe + cidx, true); // 打开中断遮罩

        lock_init(&ctrl->lock);

        for (size_t didx = 0; didx < IDE_CTRL_DISK_NR; didx++)
        {
            ide_disk_t *disk = &ctrl->devices[didx];
            disk->ctrl = ctrl;
            disk->selector = (0b1110 | didx) << 4;

            if (!ide_identify(disk))
                continue;
        }
    }
}

void ide_device_install()
{
    handler_t write = NULL;
    handler_t read = NULL;

    for (size_t cidx = 0; cidx < IDE_CTRL_NR; cidx++)
    {
        ide_ctrl_t *ctrl = &ide_ctrls[cidx];

        for (size_t didx = 0; didx < IDE_CTRL_DISK_NR; didx++)
        {
            ide_disk_t *disk = &ctrl->devices[didx];
            if (!disk->ctrl)
                continue;

            switch (ctrl->iotype)
            {
            case IDE_IO_TYPE_PIO:
                write = ide_pio_write;
                read = ide_pio_read;
                break;
            case IDE_IO_TYPE_UDMA:
                write = ide_udma_write;
                read = ide_udma_read;
                break;
            }

            disk->dev = device_create(
                DEV_TYPE_BLOCK, DEV_FLAG_DISK, -1, NULL, disk,
                 write, read, ide_ioctl);

            part_scan(disk->dev, disk->dev);
        }
    }
}

void ide_init()
{
    ide_device_init();
    ide_device_install();
}